#include "zus_wave.hpp"

constexpr uint32_t WAVE_WIDTH      = 1920;
constexpr uint32_t WAVE_HEIGHT     = 1080;
constexpr uint32_t CHANNEL_CNT     = 4;
constexpr uint32_t COLOR_NUM       = 256;
constexpr uint32_t FRAME_CNT       = 120;
constexpr uint32_t FRAME_PER_CYCLE = 90;
constexpr float    WAVE_CYCLE      = 1.5f;
constexpr float    PI              = 3.141593f;

#define CH1_COLOR 255,0,0
#define CH2_COLOR 255,255,0
#define CH3_COLOR 255,0,255
#define CH4_COLOR 0,255,255

static void linear_lut(uint8_t *lut, uint8_t r, uint8_t g, uint8_t b)
{
    auto color = reinterpret_cast<uint8_t(*)[4]>(lut);
    for (uint32_t i = 0; i<COLOR_NUM; ++i) {
        color[i][0] = i * r / 255;
        color[i][1] = i * g / 255;
        color[i][2] = i * b / 255;
        color[i][3] = 255;
    }
}

#define U8_BIT(x) ((uint8_t)(1u << (x)))

static void wave_frame(uint8_t *gray, uint8_t ch, float phase, uint32_t frame)
{
    auto rg = reinterpret_cast<uint8_t(*)[WAVE_WIDTH][2]>(gray);

    constexpr float radius_step = (2.f * PI * WAVE_CYCLE) / (float)(WAVE_WIDTH - 1);
    constexpr uint32_t h_center = WAVE_HEIGHT / 2;
    constexpr uint32_t amp = WAVE_HEIGHT / 3;
    phase += ((float)frame / FRAME_PER_CYCLE) * (2.f * PI * WAVE_CYCLE);

    for (uint32_t i=0; i<WAVE_WIDTH; ++i) {
        const float value = (float)amp * sin(phase + (float)i * radius_step);
        const uint32_t row = h_center + (uint32_t)value;

        rg[row][i][0] |= U8_BIT(ch);
        rg[row][i][1] = 255;
        for (uint32_t j=1; j<8; ++j) {
            if (row - j >= 0) {
                rg[row - j][i][0] |= U8_BIT(ch);
                rg[row - j][i][1] = std::max<uint8_t>(255 - 16*j, rg[row - j][i][1]);
            }
            if (row + j < WAVE_HEIGHT) {
                rg[row + j][i][0] |= U8_BIT(ch);
                rg[row + j][i][1] = std::max<uint8_t>(255 - 16*j, rg[row + j][i][1]);
            }
        }
    }
}

static void block_linear(uint8_t *frame, uint8_t ch, uint32_t start_row, uint32_t end_row) {
    auto rg = reinterpret_cast<uint8_t(*)[WAVE_WIDTH][2]>(frame);

    for (uint32_t row = start_row; row != end_row; ++row) {
        for (uint32_t col = 0; col < WAVE_WIDTH; ++col) {
            rg[row][col][0] |= U8_BIT(ch);
            rg[row][col][1] = col * 255 / WAVE_WIDTH;
        }
    }
}

/*
 * 2---------3
 * |  \      |
 * |      \  |
 * 0----4----1
 * |\   |\   |
 * |   \|   \|
 * 5----6----7
 * p0: -1.f, 0.f
 * p1:  1.f, 0.f
 * p2: -1.f, 1.f
 * p3:  1.f, 1.f
 * p4:  0.f, 0.f
 * p5: -1.f,-1.f
 * p6:  0.f,-1.f
 * p7: -1.f, 1.f
 */

#define p0 -1.f, 0.f
#define p1  1.f, 0.f
#define p2 -1.f, 1.f
#define p3  1.f, 1.f
#define p4  0.f, 0.f
#define p5 -1.f,-1.f
#define p6  0.f,-1.f
#define p7  1.f,-1.f
#define uv_tl 0.f, 0.f
#define uv_tr 1.f, 0.f
#define uv_bl 0.f, 1.f
#define uv_br 1.f, 1.f
#define TOP_FLAG   0b001
#define ZOOM1_FLAG 0b010
#define ZOOM2_FLAG 0b100
#define ZOOM_MASK  0b110

static float draw_data[][4] {
    // pos        // uv
    {p0,          uv_bl},
    {p1,          uv_br},
    {p2,          uv_tl},

    {p2,          uv_tl},
    {p3,          uv_tr},
    {p1,          uv_br},   // top-half end

    {p5,          uv_bl},
    {p6,          uv_br},
    {p0,          uv_tl},

    {p6,          uv_br},
    {p0,          uv_tl},
    {p4,          uv_tr},   // zoom 1 end

    {p6,          uv_bl},
    {p7,          uv_br},
    {p4,          uv_tl},

    {p4,          uv_tl},
    {p1,          uv_tr},
    {p7,          uv_br},   // zoom 2 end

    {p5,          uv_bl},
    {p7,          uv_br},
    {p2,          uv_tl},

    {p2,          uv_tl},
    {p3,          uv_tr},
    {p7,          uv_br},   // full end
};

void ZUSWave::RenderFrame(double delta_time, double total_time) {

    // upload tex
    m_tex_gray_->Upload(WAVE_WIDTH, WAVE_HEIGHT, GL_RG, GL_UNSIGNED_BYTE, m_gray_data_ + m_current_frame_ * WAVE_WIDTH * WAVE_HEIGHT * 2);

    m_prog_->Use();
    m_draw_vao_->Bind();

    m_prog_->Uniform1i("texGray", 0);
    m_prog_->Uniform1i("texLut", 1);
    m_prog_->Uniform1i("chsMask", m_enabled_channels_);

    GLCallValidation(glActiveTexture(GL_TEXTURE0));
    GLCallValidation(glBindTexture(GL_TEXTURE_2D, (GLuint)*m_tex_gray_));
    GLCallValidation(glActiveTexture(GL_TEXTURE1));
    GLCallValidation(glBindTexture(GL_TEXTURE_2D, (GLuint)*m_tex_lut_));

//    GLCallValidation(glDrawArrays(GL_TRIANGLES, 0, 18));

    if (m_area_flags_ & ZOOM_MASK) {
        // zoom enabled
        if (m_area_flags_ & TOP_FLAG)
            GLCallValidation(glDrawArrays(GL_TRIANGLES, 0, 6));
        if (m_area_flags_ & ZOOM1_FLAG)
            GLCallValidation(glDrawArrays(GL_TRIANGLES, 6, 6));
        if (m_area_flags_ & ZOOM2_FLAG)
            GLCallValidation(glDrawArrays(GL_TRIANGLES, 12, 6));
    } else if (m_area_flags_ & TOP_FLAG) {
        // no zooms
        GLCallValidation(glDrawArrays(GL_TRIANGLES, 18, 6));
    }

    m_draw_vao_->UnBind();
    m_prog_->NotUse();

    if (m_enable_animate_) {
        m_current_frame_ = static_cast<int>((total_time / (1.0 / FRAME_PER_CYCLE)));
        m_current_frame_ %= static_cast<int>(FRAME_CNT);
    }
}

void ZUSWave::RenderConfigUI() {
    using namespace ImGui;

    App::RenderConfigUI();

    if (TreeNodeEx("波形设置", ImGuiTreeNodeFlags_Framed))
    {
        AlignTextToFramePadding();
        TextUnformatted("当前帧：");
        SameLine();
        SetNextItemWidth(GetFontSize()*10);
        SliderInt("##CurFrame", &m_current_frame_, 0, FRAME_CNT-1);
        TextUnformatted("启用动画：");
        SameLine();
        Checkbox("##Animated", &m_enable_animate_);

        TextUnformatted("通道启用：");
        TextColored(ImColor(CH1_COLOR),"CH1");
        SameLine();
        CheckboxFlags("##CH1", &m_enabled_channels_, 0x1);
        TextColored(ImColor(CH2_COLOR),"CH2");
        SameLine();
        CheckboxFlags("##CH2", &m_enabled_channels_, 0x2);
        TextColored(ImColor(CH3_COLOR),"CH3");
        SameLine();
        CheckboxFlags("##CH3", &m_enabled_channels_, 0x4);
        TextColored(ImColor(CH4_COLOR),"CH4");
        SameLine();
        CheckboxFlags("##CH4", &m_enabled_channels_, 0x8);

        TextUnformatted("绘制区域：");
        CheckboxFlags("主波形(上方)", &m_area_flags_, 0b001);
//        SameLine();
        CheckboxFlags("ZOOM1(左下)", &m_area_flags_, 0b010);
//        SameLine();
        CheckboxFlags("ZOOM2(右下)", &m_area_flags_, 0b100);

        TreePop();
    }
}

void ZUSWave::OnInitialize() {
    App::OnInitialize();

    m_current_frame_ = 0;
    m_enable_animate_ = true;
    m_enabled_channels_ = 0b1111;
    m_area_flags_ = 0b001;

    // alloc memory
    m_color_map_ = new uint8_t[(CHANNEL_CNT+2)*COLOR_NUM*4];
    m_gray_data_ = new uint8_t[FRAME_CNT*WAVE_WIDTH*WAVE_HEIGHT*2];

    memset(m_color_map_, 0x00, (CHANNEL_CNT+2)*COLOR_NUM*4);
    memset(m_gray_data_, 0x00, FRAME_CNT*WAVE_WIDTH*WAVE_HEIGHT*2);

    // fill data
    linear_lut(m_color_map_ + 1 * COLOR_NUM * 4, CH1_COLOR);
    linear_lut(m_color_map_ + 2 * COLOR_NUM * 4, CH2_COLOR);
    linear_lut(m_color_map_ + 3 * COLOR_NUM * 4, CH3_COLOR);
    linear_lut(m_color_map_ + 4 * COLOR_NUM * 4, CH4_COLOR);

    // gen frames
    for (uint32_t i = 0; i<FRAME_CNT; ++i) {
        for (uint32_t j=0; j<CHANNEL_CNT; ++j) {
            wave_frame(m_gray_data_ + i * WAVE_HEIGHT * WAVE_WIDTH * 2, j+1, (float)j * 0.5f * PI, i);
            constexpr uint32_t height = WAVE_HEIGHT / 25;
            block_linear(m_gray_data_ + i * WAVE_HEIGHT * WAVE_WIDTH * 2, j + 1, j * height, height + j * height);
        }
    }

    // opengl stuff
    m_prog_ = GL::Program::Create("zus");
    m_draw_vao_ = GL::VertexArray::Create(1);
    m_draw_vao_->Bind();
    {
        auto buff = m_draw_vao_->Get(0);
        buff->Bind(GL_ARRAY_BUFFER);
        buff->Data(GL_ARRAY_BUFFER, sizeof(draw_data), draw_data, GL_STATIC_DRAW);
        GLCallValidation(glVertexAttribPointer(0, 2, GL_FLOAT, GL_FALSE, sizeof(float)*4, (void*)0));
        GLCallValidation(glVertexAttribPointer(1, 2, GL_FLOAT, GL_FALSE, sizeof(float)*4, (void*)(sizeof(float)*2)));
        GLCallValidation(glEnableVertexAttribArray(0));
        GLCallValidation(glEnableVertexAttribArray(1));
    }
    m_draw_vao_->UnBind();

    m_tex_lut_ = GL::Texture::Create();
    m_tex_lut_->Upload(COLOR_NUM, CHANNEL_CNT+2, GL_RGBA, GL_UNSIGNED_BYTE, m_color_map_);

    m_tex_gray_ = GL::Texture::Create();
}
void ZUSWave::OnCleanup() {
    App::OnCleanup();
}

DEFINE_FACTORY_FUNC(CreateZUSWaveApp, win, args) {
    (void) args;
    return std::make_unique<ZUSWave>(win, "ZUS Wave App");
}

REGISTER_APP_CREATOR("ZUSWaveApp", &CreateZUSWaveApp);
